#ifndef UnstructMG_BOX_ILU_HPP
#define UnstructMG_BOX_ILU_HPP

#include "par_CSRMatrix.hpp"

template<typename idx_t>
struct SuperTaskDep {
    idx_t dep_lev = -1;// 依赖的层号
    int   dep_tid = -1;// 依赖的线程号
    bool operator < (const SuperTaskDep & b) const {
        if (dep_lev == b.dep_lev) return dep_tid < b.dep_tid;
        else return dep_lev < b.dep_lev;
    }
};

template<typename idx_t>
struct TaskDepGraph {
    idx_t num_levs = -1;// 层数，类似于CSR里的行数

    idx_t num_deps = -1;// 依赖数，类似于CSR里的非零元数
    idx_t * dep_rpt = nullptr;// 长度 num_levs+1，dep_rpt[ilev+1] - dep_rpt[ilev] 为本线程在第ilev层的super-task的依赖数目
    SuperTaskDep<idx_t> * deps = nullptr;// 长度 dep_rpt[num_levs]，[dep_rpt[ilev], dep_rpt[ilev+1]) 为本线程在第ilev层的super-task的依赖
    
    idx_t num_rids = -1;
    idx_t * rid_rpt = nullptr;// 长度 num_levs+1，rid_rpt[ilev+1] - rid_rpt[ilev] 为本线程在第ilev层的任务的数目
    idx_t * rids = nullptr;// 长度 rid_rpt[num_levs]，[rid_rpt[ilev], rid_rpt[ilev+1]) 为本线程在第ilev层的任务的行号
    idx_t reor_off = -1;// reordered offset
    ~TaskDepGraph() {
        delete [] dep_rpt;
        delete [] deps;
        delete [] rid_rpt;
        delete [] rids;
    }
};

template<typename idx_t, typename data_t, typename setup_t, typename calc_t, int dof>
class BoxILU final : public Solver<idx_t, setup_t, calc_t, dof>
{
public:
    BoxILU(idx_t _type=0);
    ~BoxILU();
    void SetOperator(const Operator<idx_t, setup_t, dof> & op);
    void SetExecNumThreads(int num_threads);
// private:
    idx_t ilu_type = 0;
    idx_t maxfil = 100;
    setup_t droptol = 5.0e-3;
    bool own_A = false;
    const par_CSRMatrix<idx_t, data_t, calc_t, dof> * A_calc = nullptr;// setup时需要截断的低精度矩阵
    int nthreads = 1;
    TaskDepGraph<idx_t> * TDG_Ls = nullptr;
    TaskDepGraph<idx_t> * TDG_Us = nullptr;
    CSRMatrix<idx_t, data_t, calc_t, dof> L, U;
    data_t * D_inv = nullptr;
    calc_t * aux_array = nullptr;
    void factorize_lev0(const CSRMatrix<idx_t, setup_t, setup_t, dof> & A);
    void factorize_threshold(const CSRMatrix<idx_t, setup_t, setup_t, dof> & A);
    void forward (const calc_t * rhs, calc_t * sol) const;
    void backward(const calc_t * rhs, calc_t * sol) const;

protected:
    TaskDepGraph<idx_t> * build_TDG(CSRMatrix<idx_t, data_t, calc_t, dof> & tri_mat, const char tri_type);
    void truncate();
    void Mult_ipl(const par_Vector<idx_t, calc_t, dof> & input, 
                        par_Vector<idx_t, calc_t, dof> & output) const ;
};
typedef BoxILU<int, double, double, double, 1> BoxILU_I32All64_DOF1;
typedef BoxILU<int, double, double, double, 2> BoxILU_I32All64_DOF2;
typedef BoxILU<int, double, double, double, 3> BoxILU_I32All64_DOF3;
typedef BoxILU<int, double, double, double, 4> BoxILU_I32All64_DOF4;

typedef BoxILU<int, float , double, double, 1> BoxILU_I32K64P64D32_DOF1;
typedef BoxILU<int, float , double, double, 2> BoxILU_I32K64P64D32_DOF2;
typedef BoxILU<int, float , double, double, 3> BoxILU_I32K64P64D32_DOF3;
typedef BoxILU<int, float , double, double, 4> BoxILU_I32K64P64D32_DOF4;

typedef BoxILU<int, float , double, float , 1> BoxILU_I32K64P32D32_DOF1;
typedef BoxILU<int, float , double, float , 2> BoxILU_I32K64P32D32_DOF2;
typedef BoxILU<int, float , double, float , 3> BoxILU_I32K64P32D32_DOF3;
typedef BoxILU<int, float , double, float , 4> BoxILU_I32K64P32D32_DOF4;

typedef BoxILU<int, float , float , float , 1> BoxILU_I32All32_DOF1;
typedef BoxILU<int, float , float , float , 2> BoxILU_I32All32_DOF2;
typedef BoxILU<int, float , float , float , 3> BoxILU_I32All32_DOF3;
typedef BoxILU<int, float , float , float , 4> BoxILU_I32All32_DOF4;
#ifdef USE_FP16
typedef BoxILU<int, __fp16, double, float , 1> BoxILU_I32K64P32D16_DOF1;
typedef BoxILU<int, __fp16, double, float , 2> BoxILU_I32K64P32D16_DOF2;
typedef BoxILU<int, __fp16, double, float , 3> BoxILU_I32K64P32D16_DOF3;
typedef BoxILU<int, __fp16, double, float , 4> BoxILU_I32K64P32D16_DOF4;

typedef BoxILU<int, __fp16, float , float , 1> BoxILU_I32K32P32D16_DOF1;
typedef BoxILU<int, __fp16, float , float , 2> BoxILU_I32K32P32D16_DOF2;
typedef BoxILU<int, __fp16, float , float , 3> BoxILU_I32K32P32D16_DOF3;
typedef BoxILU<int, __fp16, float , float , 4> BoxILU_I32K32P32D16_DOF4;
#endif

#endif